/**
 * Created by frankqian on 2018/7/22.
 */

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.*;

import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.StringTokenizer;

import org.apache.hadoop.mapreduce.*;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import org.apache.hadoop.util.GenericOptionsParser;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import randomdata.FormatData;

//import org.apache.hadoop.mapreduce.Job;
//import org.apache.hadoop.mapreduce.JobContext;


class GenDataMapper extends Mapper<LongWritable, NullWritable, Text, String> {
//    Text datarow = new Text();
//private static final Logger LOG = LoggerFactory.getLogger(GenDataMapper.class);
    public void map(LongWritable row, NullWritable ignored,
                    Context context) throws IOException,InterruptedException {

//            datarow.set(itr.nextToken());
        String record = FormatData.getARecord();
        context.write(new Text(row.toString()), record);
    }
}

public class GenData {

    static final String NUM_ROWS_KEY = "num_of_row";

    private static final Logger LOG = LoggerFactory.getLogger(GenData.class);

    static class RangeInputFormat
            extends InputFormat<LongWritable, NullWritable> {

        /**
         * An input split consisting of a range on numbers.
         */
        static class RangeInputSplit extends InputSplit implements Writable {
            long firstRow;
            long rowCount;

            public RangeInputSplit() {
            }

            public RangeInputSplit(long offset, long length) {
                firstRow = offset;
                rowCount = length;
            }

            public long getLength() throws IOException {
                return 0;
            }

            public String[] getLocations() throws IOException {
                return new String[]{};
            }

            public void readFields(DataInput in) throws IOException {
                firstRow = WritableUtils.readVLong(in);
                rowCount = WritableUtils.readVLong(in);
            }

            public void write(DataOutput out) throws IOException {
                WritableUtils.writeVLong(out, firstRow);
                WritableUtils.writeVLong(out, rowCount);
            }
        }

        /**
         * Create the desired number of splits, dividing the number of rows
         * between the mappers.
         */
        public List<InputSplit> getSplits(JobContext job) {
            long totalRows = getNumberOfRows(job);
            int numSplits = job.getConfiguration().getInt(MRJobConfig.NUM_MAPS, 1);
            LOG.info("Generating " + totalRows + " using " + numSplits);
            List<InputSplit> splits = new ArrayList<InputSplit>();
            long currentRow = 0;
            for(int split = 0; split < numSplits; ++split) {
                long goal =
                        (long) Math.ceil(totalRows * (double)(split + 1) / numSplits);
                splits.add(new RangeInputSplit(currentRow, goal - currentRow));
                currentRow = goal;
            }
            return splits;
        }

        /**
         * A record reader that will generate a range of numbers.
         */
        static class RangeRecordReader
                extends RecordReader<LongWritable, NullWritable> {
            long startRow;
            long finishedRows;
            long totalRows;
            LongWritable key = null;

            public RangeRecordReader() {
            }

            public void initialize(InputSplit split, TaskAttemptContext context)
                    throws IOException, InterruptedException {
                startRow = ((RangeInputSplit)split).firstRow;
                finishedRows = 0;
                totalRows = ((RangeInputSplit)split).rowCount;
            }

            public void close() throws IOException {
                // NOTHING
            }

            public LongWritable getCurrentKey() {
                return key;
            }

            public NullWritable getCurrentValue() {
                return NullWritable.get();
            }

            public float getProgress() throws IOException {
                return finishedRows / (float) totalRows;
            }

            public boolean nextKeyValue() {
                if (key == null) {
                    key = new LongWritable();
                }
                if (finishedRows < totalRows) {
                    key.set(startRow + finishedRows);
                    finishedRows += 1;
                    return true;
                } else {
                    return false;
                }
            }

        }

        public RecordReader<LongWritable, NullWritable>
        createRecordReader(InputSplit split, TaskAttemptContext context)
                throws IOException {
            return new RangeRecordReader();
        }

    }

    static long getNumberOfRows(JobContext job) {
        return job.getConfiguration().getLong(NUM_ROWS_KEY,
                1);
    }

    static void setNumberOfRows(Job job, long numRows) {
        job.getConfiguration().setLong(NUM_ROWS_KEY, numRows);
    }

    private static long parseHumanLong(String str) {
        long num = Integer.parseInt(str.substring(0, str.length() - 1));
        switch (str.substring(str.length() - 1)) {
            case "t":
            case "T":
                num =  num * 1024 * 1024 * 1024 * 1024;
                break;
            case "g":
            case "G":
                num =  num * 1024 * 1024 * 1024;
                break;
            case "m":
            case "M":
                num =  num * 1024 * 1024;
                break;
            case "k":
            case "K":
                num =  num * 1024;
                break;
            case "b":
            case "B":
                num =  num;
                break;
            default:
                throw new IllegalArgumentException("wrong data size");
        }
        long batchsize =  num / FormatData.getARecord().length();
        return batchsize;
    }

    public static void main(String[] args) throws IOException, ClassNotFoundException, InterruptedException {
        Configuration conf = new Configuration();
        String[] otherArgs = new GenericOptionsParser(conf, args).getRemainingArgs();
        if(otherArgs.length != 2) {
            System.err.println("gendata <num of data> <out path>");
            System.exit(2);
        }
        conf.setInt(MRJobConfig.NUM_MAPS, 8);
        Job job = Job.getInstance(conf, "gendata");
        setNumberOfRows(job, parseHumanLong(args[0]));
        job.setJarByClass(GenData.class);
        job.setMapperClass(GenDataMapper.class);
        job.setInputFormatClass(RangeInputFormat.class);
        job.setOutputKeyClass(Text.class);
        job.setOutputValueClass(String.class);
        job.setNumReduceTasks(0);
        FileOutputFormat.setOutputPath(job, new Path(otherArgs[1]));
        System.exit(job.waitForCompletion(true)?0:1);

    }
}

